import elasticsearch
import openpyxl
import os
import pathlib
import pandas as pd
import sys
import unittest

project_path = str(pathlib.Path(os.path.abspath(__file__)).parent.parent.parent.parent)
sys.path.append(project_path)
from shell.knowledge_graph import question_parser, question_classifier
from unittest import mock

shell = __import__("shell.import", globals(), locals(), ['settings'], 0)


class QAParserTest(unittest.TestCase):
    @classmethod
    def setUpClass(self) -> None:
        self.__test_excel_file = (pathlib.Path(os.path.abspath(__file__)).parent.parent.parent.parent /
                           'doc' /
                           '知识图谱问答系统开发计划.xlsx')
        test_sheet_name = '问题案例模板'
        columns = ['问题大类（节点）', '关系/节点', '问题子类（属性）', '详细问题', '是否聚合']
        test_df = pd.read_excel(self.__test_excel_file, sheet_name=test_sheet_name)

        self.__test_df = test_df[columns]
        self.__es = elasticsearch.Elasticsearch()
        self.__qc = None
        self.__wk_sheets = openpyxl.load_workbook(self.__test_excel_file)
        self.__sheet = self.__wk_sheets[test_sheet_name]
        
        # 初始化测试结果列为空，背景色设置为白色
        white_fill = openpyxl.styles.PatternFill("solid", fgColor='FFFFFF')
        for cell in list(self.__sheet.rows)[0]:
            if cell.value == '问题案例es测试子类结果':
                # 获取列字母(这样的-> A B C D)
                self.__test_result_col = openpyxl.utils.get_column_letter(cell.column)
                break
        for row in range(self.__sheet.max_row):
            if row == 0:
                continue
            self.__sheet[self.__test_result_col + str(row + 1)].value = ''
            self.__sheet[self.__test_result_col + str(row + 1)].fill = white_fill

    @classmethod
    def tearDownClass(self) -> None:
        self.__wk_sheets.save(self.__test_excel_file)

    def test_excelcase_attr(self) -> None:
        faker_path = pathlib.Path(os.path.abspath(__file__)).parent.parent.parent / 'test_fake_data'
        sys.path.append(str(faker_path))
        import gen_fake_data as gfd

        for i, row in self.__test_df.iterrows():
            if row[4] == '是':
                continue

            node2name = row[1].split('|')[1]
            node2cnattr = row[2].split('|')[0]
            node2attr = row[2].split('|')[1]
            attr_question = row[3].split(']')[1]

            _, attr_dict = question_parser.QAParser.es_parser(attr_question, self.__es, self.__qc)
            if not attr_dict:
                fill = openpyxl.styles.PatternFill("solid", fgColor='FF0000')
            else:
                if (attr_dict['value'] == node2cnattr and
                    attr_dict['table'] == node2name and
                    attr_dict['attr'] == node2attr):
                    fill = openpyxl.styles.PatternFill("solid", fgColor='FFFFFF')
                else:
                    fill = openpyxl.styles.PatternFill("solid", fgColor='FF0000')
                test_result = f"{attr_dict['value']}|{attr_dict['table']}|{attr_dict['attr']}"
                self.__sheet[self.__test_result_col + str(i + 2)].value = test_result

            self.__sheet[self.__test_result_col + str(i + 2)].fill = fill

    def __test_excelcase(self) -> None:
        faker_path = pathlib.Path(os.path.abspath(__file__)).parent.parent.parent / 'test_fake_data'
        sys.path.append(str(faker_path))
        import gen_fake_data as gfd

        for i, row in self.__test_df.iterrows():
            if row[4] == '是':
                continue
            # self.__es.search = mock.Mock(return_value={'hits': {'hits': [
            #                                     {'_source': {'is_attr': 'True', 'table': 'node1'}},
            #                                     {'_source': {'is_attr': 'False', 'table': 'node2'}}]}})
            value_names = row[3].split('[')[1].split(']')[0]

            for value_name in value_names.split('|'):
                node1attr = value_name.split(',')[1]
                node1cnattr = value_name.split(',')[0]
                node1name = row[0].split('|')[1]
                node1value = str(gfd.get_sample(shell.settings.nodedict[node1name], node1cnattr))
                node2name = row[1].split('|')[1]
                node2cnattr = row[2].split('|')[0]
                node2attr = row[2].split('|')[1]
                attr_result = (f"{{'value': '{node2cnattr}', 'attr': '{node2attr}', "
                               f"'table': '{node2name}', 'is_attr': 'True'}}")
                value_result = (f"{{'index': 'kgqa_value', 'value': '{node1value}', 'attr': '{node1attr}', "
                                f"'table': '{node1name}', 'is_attr': 'False'}}")

                question = node1value + row[3].split(']')[1]
                testitem = (f"line {i + 2}, "
                            f"node1={{ {row[0]}, {value_name}, {node1value}}}, "
                            f"node2={{ {row[1]}, {row[2]} }}, "
                            f"question={question}")
                with self.subTest(testitem=testitem):
                    value_dict, attr_dict = question_parser.QAParser.es_parser(question, self.__es, self.__qc)
                    self.assertEqual(str(value_dict), value_result)
                    self.assertEqual(str(attr_dict), attr_result)

    def test_sample(self) -> None:
        # self.__es.search = mock.Mock(return_value={'hits': {'hits': [
        #                                     {'_source': {'is_attr': 'True', 'table': 'node1'}},
        #                                     {'_source': {'is_attr': 'False', 'table': 'node2'}}]}})
        self.assertEqual(str(question_parser.QAParser.es_parser('whatever', self.__es, self.__qc)),
                         "({}, {})")


if __name__ == '__main__':
    suite = unittest.TestSuite()
    suite.addTests(unittest.TestLoader().loadTestsFromTestCase(QAParserTest))

    resultfile = pathlib.Path(os.path.dirname(__file__)) / (__file__.split('.')[0] + '.report')
    with open(resultfile, 'w', encoding='utf8') as f:
        runner = unittest.TextTestRunner(stream=f, verbosity=2)
        runner.run(suite)
